{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SageMaker V3 JumpStart E2E Training Example\n",
    "\n",
    "This notebook demonstrates how to use SageMaker V3 to train a JumpStart model from scratch and deploy it for inference."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prerequisites\n",
    "Note: Ensure you have sagemaker and ipywidgets installed in your environment. The ipywidgets package is required to monitor endpoint deployment progress in Jupyter notebooks.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import required libraries\n",
    "import json\n",
    "import uuid\n",
    "\n",
    "from sagemaker.serve.model_builder import ModelBuilder\n",
    "from sagemaker.train.model_trainer import ModelTrainer\n",
    "from sagemaker.core.jumpstart.configs import JumpStartConfig\n",
    "from sagemaker.core.resources import EndpointConfig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Configure JumpStart Model for Training\n",
    "\n",
    "We'll train a HuggingFace Falcon model using JumpStart."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configuration\n",
    "MODEL_ID = \"huggingface-spc-bert-base-cased\"\n",
    "MODEL_NAME_PREFIX = \"js-e2e-example-model\"\n",
    "ENDPOINT_NAME_PREFIX = \"js-e2e-example-endpoint\"\n",
    "\n",
    "# Generate unique identifiers\n",
    "unique_id = str(uuid.uuid4())[:8]\n",
    "training_job_name = f\"js-training-{unique_id}\"\n",
    "model_name = f\"{MODEL_NAME_PREFIX}-{unique_id}\"\n",
    "endpoint_name = f\"{ENDPOINT_NAME_PREFIX}-{unique_id}\"\n",
    "\n",
    "print(f\"Training job name: {training_job_name}\")\n",
    "print(f\"Model name: {model_name}\")\n",
    "print(f\"Endpoint name: {endpoint_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Train the Model\n",
    "\n",
    "Use ModelTrainer to train the JumpStart model. The training job may take 30+ minutes to complete."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create JumpStart configuration\n",
    "jumpstart_config = JumpStartConfig(model_id=MODEL_ID)\n",
    "\n",
    "# Initialize ModelTrainer from JumpStart config\n",
    "model_trainer = ModelTrainer.from_jumpstart_config(\n",
    "    jumpstart_config=jumpstart_config, \n",
    "    base_job_name=training_job_name, \n",
    "    hyperparameters={\"epochs\": 1}\n",
    ")\n",
    "\n",
    "# Train the model\n",
    "print(\"Starting model training...\")\n",
    "model_trainer.train()\n",
    "print(f\"Training completed: {training_job_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Build ModelBuilder from Trained Model\n",
    "\n",
    "Create a ModelBuilder from the training artifacts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build ModelBuilder from trained model\n",
    "model_builder = ModelBuilder(\n",
    "    model=model_trainer,\n",
    "    dependencies={\"auto\": False}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Build the Model\n",
    "\n",
    "Build the model artifacts and prepare for deployment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build the model\n",
    "core_model = model_builder.build(model_name=model_name)\n",
    "print(f\"Model Successfully Created: {core_model.model_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: Deploy the Trained Model\n",
    "\n",
    "Deploy the trained model to a SageMaker endpoint for real-time inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Deploy the trained model to an endpoint\n",
    "core_endpoint = model_builder.deploy(endpoint_name=endpoint_name)\n",
    "print(f\"Endpoint Successfully Created: {core_endpoint.endpoint_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6: Test the Endpoint\n",
    "\n",
    "This endpoint performs text entailment classification, determining the logical relationship between pairs of sentences. The returned scores indicate the model's confidence for different entailment categories (e.g., entailment, contradiction, neutral) - higher scores indicate stronger predictions for each relationship type. \n",
    "\n",
    "This sends a test request to the deployed endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test with a sample query for the trained model\n",
    "test_data = [\"The weather is sunny today\", \"It is not raining\"]\n",
    "\n",
    "result = core_endpoint.invoke(\n",
    "    body=json.dumps(test_data),\n",
    "    content_type=\"application/list-text\"\n",
    ")\n",
    "\n",
    "entailment_scores = json.loads(result.body.read().decode('utf-8'))\n",
    "print(f\"Result of invoking trained endpoint: {entailment_scores}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 7: Clean Up Resources\n",
    "\n",
    "Clean up the created resources to avoid ongoing charges."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clean up resources\n",
    "core_endpoint_config = EndpointConfig.get(endpoint_config_name=core_endpoint.endpoint_name)\n",
    "\n",
    "# Delete in the correct order\n",
    "core_model.delete()\n",
    "core_endpoint.delete()\n",
    "core_endpoint_config.delete()\n",
    "\n",
    "print(\"Trained Model and Endpoint Successfully Deleted!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This notebook demonstrated:\n",
    "1. Training a JumpStart model using ModelTrainer\n",
    "2. Building a ModelBuilder from training artifacts\n",
    "3. Building the model for deployment\n",
    "4. Deploying to a SageMaker endpoint\n",
    "5. Making inference requests\n",
    "6. Cleaning up resources\n",
    "\n",
    "The V3 ModelTrainer and ModelBuilder provide a seamless end-to-end workflow from training to deployment!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv-test",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
